import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.util.ArrayDeque;
import jdk.internal.vm.Continuation;
import jdk.internal.vm.ContinuationScope;
import util.UnsafeUtil;

// javac --add-exports java.base/jdk.internal.vm=ALL-UNNAMED Prime3a.java
// java  --add-opens   java.base/jdk.internal.vm=ALL-UNNAMED Prime3a
public class Prime3a {
	private static final long CONT_OFFSET;
	private static final MethodHandle mhYield0, mhEnterSpecial;
	private static final ContinuationScope cs = new ContinuationScope("VirtualThreads");
	private static final Continuation[] readyStack = new Continuation[10000];
	private static int readyStackSize;

	static {
		try {
			var lookup = MethodHandles.lookup();
			CONT_OFFSET = UnsafeUtil.objectFieldOffset(Thread.class, "cont");
			var m = UnsafeUtil.getMethod(Continuation.class, "yield0", ContinuationScope.class, Continuation.class);
			mhYield0 = lookup.unreflect(m);
			m = UnsafeUtil.getMethod(Continuation.class, "enterSpecial", Continuation.class, boolean.class, boolean.class);
			mhEnterSpecial = lookup.unreflect(m);
		} catch (ReflectiveOperationException e) {
			throw new RuntimeException(e);
		}
	}

	@SuppressWarnings("removal")
	static void loop() {
		try {
			var t = Thread.currentThread();
			var s = readyStack;
			for (int n; (n = readyStackSize) > 0; ) {
				var c = s[--n];
				s[n] = null;
				readyStackSize = n;
				UnsafeUtil.unsafe.putObject(t, CONT_OFFSET, c);
				mhEnterSpecial.invokeExact(c, true, false);
			}
		} catch (Throwable e) {
			throw new RuntimeException(e);
		}
	}

	static void resume(Continuation c) {
		if (c.isDone())
			throw new IllegalStateException("Continuation terminated");
		int n = readyStackSize;
		readyStack[n] = c;
		readyStackSize = n + 1;
	}

	static void go(Runnable r) {
		new Continuation(cs, r).run();
	}

	@SuppressWarnings("removal")
	static void pause(ArrayDeque<Continuation> queue) {
		var c = (Continuation)UnsafeUtil.unsafe.getObject(Thread.currentThread(), CONT_OFFSET);
		queue.addLast(c);
		try {
			@SuppressWarnings("unused")
			var __ = (boolean)mhYield0.invokeExact(c, cs, (Continuation)null);
		} catch (Throwable e) {
			throw new RuntimeException(e);
		}
	}

	static class IntChan {
		private final ArrayDeque<Continuation> waitQueue = new ArrayDeque<>();
		private int value;
		private boolean hasValue;

		void put(int v) {
			if (hasValue)
				pause(waitQueue);
			value = v;
			hasValue = true;
			var c = waitQueue.pollFirst();
			if (c != null)
				resume(c);
		}

		int get() {
			if (!hasValue)
				pause(waitQueue);
			hasValue = false;
			var c = waitQueue.pollFirst();
			if (c != null)
				resume(c);
			return value;
		}
	}

	private static void generate(IntChan ch) {
		//noinspection InfiniteLoopStatement
		for (int i = 2; ; i++)
			ch.put(i);
	}

	private static void filter(IntChan in, IntChan out, int prime) {
		//noinspection InfiniteLoopStatement
		for (; ; ) {
			var i = in.get();
			if (i % prime != 0)
				out.put(i);
		}
	}

	public static void main(String[] args) {
		var count = args.length > 0 ? Integer.parseInt(args[0]) : 10000;
		var ch = new IntChan();
		go(() -> generate(ch));
		go(() -> {
			var ch0 = ch;
			for (int i = 0; ; ) {
				var prime = ch0.get();
				if (++i == count) {
					System.out.println(prime);
					System.exit(0);
				}
				var ch1 = ch0;
				var ch2 = new IntChan();
				go(() -> filter(ch1, ch2, prime));
				ch0 = ch2;
			}
		});
		loop();
	}
}
